package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

type JudgeResult struct {
	Stdout string
	Stderr string
	ExitCode int
	Timeout bool
}

var nsjailPath = "/usr/local/bin/nsjail"  // nsjail 二进制路径

func runSandbox(command []string, timeLimit int, memoryMB int) JudgeResult {
	args := []string {
		"-Mo",					// standalone once
		"--chroot", "/",		// 根目录沙箱
		"--time_limit", fmt.Sprint(timeLimit),
		"--rlimit_as", fmt.Sprint(memoryMB),
		"--",   				// 后面是实际要执行的命令
	}
	args = append(args, command...)

	cmd := exec.Command(nsjailPath, args...)
	output, err := cmd.CombinedOutput()

	result := JudgeResult {
		Stdout: string(output),
		ExitCode: 0,
		Timeout: false,
	}

	if err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			result.ExitCode = exitError.ExitCode()
		} else {
			result.Stderr = err.Error()
		}
	}

	return result
}

func judgePython(filePath string) JudgeResult {
	tmpDir, _ := ioutil.TempDir("", "oj")
	defer os.RemoveAll(tmpDir)

	tmpFile := filepath.Join(tmpDir, "main.py")
	data, _ := ioutil.ReadFile(filePath)
	_ = ioutil.WriteFile(tmpFile, data, 0644)

	return runSandbox([]string{"/usr/bin/python3", tmpFile}, 2, 64)
}

func judgeC(filePath string) JudgeResult {
	tmpDir, _ := ioutil.TempDir("", "oj")
	defer os.RemoveAll(tmpDir)

	exePath := filepath.Join(tmpDir, "a.out")

	// 编译 C 代码
	build := exec.Command("gcc", filePath, "-o", exePath)
	buildOutput, err := build.CombinedOutput()
	if err != nil {
		return JudgeResult{
			Stdout:   string(buildOutput),
			Stderr:   err.Error(),
			ExitCode: build.ProcessState.ExitCode(),
		}
	}

	return runSandbox([]string{exePath}, 2, 64)
}

func main() {
	submissions := []string{
		"./submissions/test.py",
		"./submissions/test.c",
	}

	for _, file := range submissions {
		fmt.Println("Judging:", file)
		ext := filepath.Ext(file)
		var res JudgeResult

		switch ext {
		case ".py":
			res = judgePython(file)
		case ".c":
			res = judgeC(file)
		default:
			fmt.Println("Unsupported language:", ext)
			continue
		}

		fmt.Printf("ExitCode: %d\n", res.ExitCode)
		fmt.Printf("Stdout:\n%s\n", res.Stdout)
		fmt.Printf("Stderr:\n%s\n", res.Stderr)
		fmt.Println(strings.Repeat("-", 30))
		time.Sleep(500 * time.Millisecond) // 防止同时执行过多
	}
}